#!/bin/bash
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/alias.sh"
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/conda_run.sh"
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/conda_and_pip_settings.sh"
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/print_sys_info.sh"

APP_INSTALL_NAME="fireredtts2"
APP_TARGET_DIR_NAME="fireredtts2"

export CURRENT_APP_PATH="$LMD_APPS_DIR/$APP_TARGET_DIR_NAME"
echo "App Dir: ${CURRENT_APP_PATH}"

cd "$CURRENT_APP_PATH"

if [ -z "$HF_ENDPOINT" ] && [ -n "$HF_MIRROR" ]; then
    export HF_ENDPOINT="$HF_MIRROR"
fi

export GRADIO_TEMP_DIR="$CURRENT_APP_PATH/temp"
CONDA_PREFIX="./venv"

echo "use ./venv to start app."
export HF_HOME="$CURRENT_APP_PATH/tmp_hf_home"
export MODELSCOPE_HOME="$CURRENT_APP_PATH/tmp_modelscope_home"
export MODELSCOPE_CACHE="$CURRENT_APP_PATH/tmp_modelscope_cache"

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True


LMD_CURRENT_DEVICE="cpu"
if [ -f /System/Library/CoreServices/SystemVersion.plist ]; then
  echo "macOS system."
  ARCH=$(uname -m)
  if [ "$ARCH" == "x86_64" ]; then
      echo "Intel-based Mac. device: cpu"
      LMD_CURRENT_DEVICE="cpu"
  else
      echo "Apple Silicon (M-series) Mac. device: mps"
      LMD_CURRENT_DEVICE="mps"
  fi
elif current_os_is_windows; then
  if [ "$LMD_LAUNCH_APP_MODULE" = "cpu-mode" ]; then
    export CUDA_VISIBLE_DEVICES="-1"
    LMD_CURRENT_DEVICE="cpu"
  else
    echo "Check NVIDIA gpu."
    if ! command -v nvidia-smi &> /dev/null; then
      echo "No NVIDIA driver detected."
    else
      LMD_CURRENT_DEVICE="cuda"
      echo "NVIDIA driver detected."
    fi
  fi
fi

export TORCH_LOAD_MAP_LOCATION="$LMD_CURRENT_DEVICE"
echo "TORCH_LOAD_MAP_LOCATION=$TORCH_LOAD_MAP_LOCATION"

cp "$LMD_BASE_INSTALL_SCRIPT_DIR/apps/$APP_INSTALL_NAME/patch_torch.py" "$CURRENT_APP_PATH"
cp "$LMD_BASE_INSTALL_SCRIPT_DIR/apps/$APP_INSTALL_NAME/lmd_run.py" "$CURRENT_APP_PATH"

# replace device and add device param
# sed "s/from fireredtts2.fireredtts2/from fireredtts2.fireredtts2_temp/g" "gradio_demo.py" > "gradio_demo_temp1.py"

sed "s/initiate_model(pretrained_dir: str, device=\"cuda\")/initiate_model(pretrained_dir: str, device=\"${LMD_CURRENT_DEVICE}\")/g" "gradio_demo.py" > "gradio_demo_temp1.py"

# sed "s/RedCodecInfer.from_pretrained(codec_config_path, codec_ckpt_path)/RedCodecInfer.from_pretrained(codec_config_path, codec_ckpt_path, device=device)/g" "fireredtts2/fireredtts2.py" > "fireredtts2/fireredtts2_temp.py"

conda run -p $CONDA_PREFIX --no-capture-output python "lmd_run.py"  --pretrained-dir "./pretrained_models/FireRedTTS2"
